"""功耗采样器 - 基于后端进行功耗数据采集"""

from __future__ import annotations

import time
from contextlib import contextmanager
from typing import List, Tuple, Optional

from .power_backends import get_power_backend, PowerBackend


class PowerSampler:
    """功耗采样器 - 使用指定后端进行功耗采样"""
    
    def __init__(
        self,
        backend_name: Optional[str] = None,
        device_id: int = 0,
        sample_rate_hz: float = 50.0
    ):
        """
        初始化功耗采样器
        
        Args:
            backend_name: 功耗后端名称（如 "nvml", "rocm", "rapl"）
                        如果为 None，将自动检测可用后端
            device_id: 设备ID
            sample_rate_hz: 采样频率（Hz）
        """
        self.device_id = device_id
        self.sample_rate_hz = sample_rate_hz
        self.sample_interval = 1.0 / sample_rate_hz if sample_rate_hz > 0 else 0.1
        
        # 创建后端实例
        if backend_name:
            self.backend = get_power_backend(backend_name)
        else:
            # 自动检测可用后端
            self.backend = self._auto_detect_backend()
        
        # 初始化后端
        if not self.backend.init(device_id):
            raise RuntimeError(
                f"Failed to initialize power backend '{self.backend.get_name()}' "
                f"for device {device_id}"
            )
        
        # 采样数据存储：[(timestamp_ms, power_mW), ...]
        self.samples: List[Tuple[float, float]] = []
        self._sampling_active = False
    
    def _auto_detect_backend(self) -> PowerBackend:
        """
        自动检测可用的功耗后端
        
        Returns:
            PowerBackend 实例
            
        Raises:
            RuntimeError: 没有可用的后端
        """
        from .power_backends import list_power_backends
        
        available = list_power_backends(include_unavailable=False)
        if not available:
            raise RuntimeError(
                "No power backend available. Please install required dependencies:\n"
                "  - NVIDIA GPU: pip install pynvml\n"
                "  - AMD GPU: Install rocm-smi\n"
                "  - Intel CPU: Ensure RAPL interface is available"
            )
        
        # 按优先级选择后端
        priority_order = ["nvml", "rocm", "rapl", "ascend", "external_meter"]
        for name in priority_order:
            if name in available:
                return get_power_backend(name)
        
        # 如果优先级中没有，返回第一个可用
        return get_power_backend(available[0])
    
    @contextmanager
    def sampling(self):
        """
        上下文管理器：在进入时开始采样，退出时停止采样
        
        使用示例:
            with power_sampler.sampling():
                # 执行需要监控的代码
                run_inference()
            # 采样数据存储在 power_sampler.samples 中
        """
        self.start()
        try:
            yield
        finally:
            self.stop()
    
    def start(self):
        """开始采样（清空之前的采样数据）"""
        self.samples.clear()
        self._sampling_active = True
        
        # 在后台线程中进行采样
        import threading
        self._sampling_thread = threading.Thread(target=self._sampling_loop, daemon=True)
        self._sampling_thread.start()
    
    def stop(self):
        """停止采样"""
        self._sampling_active = False
        if hasattr(self, '_sampling_thread'):
            self._sampling_thread.join(timeout=2.0)
    
    def _sampling_loop(self):
        """采样循环（在后台线程中运行）"""
        start_time = time.perf_counter()
        
        while self._sampling_active:
            # 采样当前功率
            power_mW = self.backend.sample()
            
            if power_mW is not None:
                # 计算相对时间戳（毫秒）
                elapsed_ms = (time.perf_counter() - start_time) * 1000.0
                self.samples.append((elapsed_ms, power_mW))
            
            # 等待下一个采样周期
            time.sleep(self.sample_interval)
    
    def cleanup(self):
        """清理资源"""
        self.stop()
        if hasattr(self, 'backend'):
            self.backend.cleanup()
    
    def __del__(self):
        """析构函数：确保资源被清理"""
        try:
            self.cleanup()
        except:
            pass

